import json
import re
from rag_chain import RAGChain
from text_processor import TextVectorizer

# 目标电话号码和公司
target_phone = "15937098754"
expected_company = "北京科技创新有限公司"

# 测试数据
test_data = """
这是一个测试文件，包含电话号码与公司的关联信息。

以下是一些公司和联系人信息：
1. 北京科技创新有限公司 - 张三 - 15937098754
2. 上海未来信息技术有限公司 - 李四 - 13812345678
3. 广州智慧城市发展有限公司 - 王五 - 13987654321
"""

# 初始化测试
print("=== 开始详细RAG链测试 ===")

# 1. 初始化向量器和RAG链
sim_metric = "cosine"
vectorizer = TextVectorizer(use_tfidf=True)
rag_chain = RAGChain(deepseek_api_key="sk-a6d2da9cbeaf4cf59512a993de770cbd", vectorizer=vectorizer)
print(f"已初始化向量器和RAG链，相似度度量: {sim_metric}")

# 2. 手动添加测试数据到向量器
vectorizer.vectorize_and_store(test_data, metadata={"source": "test_phone_company.txt"})
print("已手动添加测试数据到向量器")

# 3. 检查词汇表
print("\n=== 词汇表检查 ===")
print(f"词汇表大小: {len(vectorizer.vocab)}")
print(f"电话号码 '{target_phone}' 在词汇表中: {target_phone in vectorizer.vocab}")
print(f"公司名称 '{expected_company}' 在词汇表中: {expected_company in vectorizer.vocab}")

# 4. 直接使用向量器查询
print("\n=== 向量器直接查询 ===")
vec_results = vectorizer.query(target_phone, n_results=3, similarity_metric=sim_metric)
print(f"查询结果: {json.dumps(vec_results, ensure_ascii=False, indent=2)}")

# 5. 使用RAG链的retrieve_relevant_documents方法
print("\n=== RAG链文档检索 ===")
rag_docs = rag_chain.retrieve_relevant_documents(target_phone, n_results=3, similarity_metric=sim_metric)
print(f"检索到的文档: {json.dumps(rag_docs, ensure_ascii=False, indent=2)}")

# 6. 构建提示词
print("\n=== 构建提示词 ===")
prompt = rag_chain.build_prompt(target_phone, rag_docs)
print(f"提示词: {prompt}")

# 7. 检查查询扩展
print("\n=== 查询扩展检查 ===")
expanded_query = rag_chain.expand_query(target_phone)
print(f"原始查询: {target_phone}")
print(f"扩展查询: {expanded_query}")

# 8. 验证结果
print("\n=== 结果验证 ===")
found_in_vec = any(target_phone in doc.get("text", "") and expected_company in doc.get("text", "") for doc in vec_results)
found_in_rag = any(target_phone in doc.get("text", "") and expected_company in doc.get("text", "") for doc in rag_docs)

print(f"向量器查询中找到关联信息: {found_in_vec}")
print(f"RAG链检索中找到关联信息: {found_in_rag}")

if found_in_vec and not found_in_rag:
    print("问题可能出在RAG链的retrieve_relevant_documents方法中")
elif found_in_rag:
    print("RAG链检索到了关联信息，但AI回答可能未正确提取，可能需要优化提示模板")
else:
    print("向量器和RAG链都未找到关联信息，可能是向量器配置问题")

print("\n=== 测试完成 ===")